import os
import json
import random
import math
import importlib
import numpy as np
import warnings

warnings.filterwarnings('ignore', category=FutureWarning)
from shapely.geometry import Point, Polygon
import matplotlib.pyplot as plt

from hybrid_a_star.hybrid_a_star import pi_2_pi

# 导入onsite-mine相关模块
## onsite-mine.dynamic_scenes模块
scenarioOrganizer = importlib.import_module("onsite-unstructured.dynamic_scenes.scenarioOrganizer1")
env = importlib.import_module("onsite-unstructured.dynamic_scenes.env")
controller = importlib.import_module("onsite-unstructured.dynamic_scenes.controller")
lookup = importlib.import_module("onsite-unstructured.dynamic_scenes.lookup")


# grid_size = 600  # 网格大小
# real_step_size = 5  # 寻路步长
lim_angle = 30  # 默认最大偏航角
dynamic_plot = False  # 动态画图
do_plot = False

random.seed(0)  # 设置随机种子

# 检查对应文件夹是否存在
def check_dir(target_dir):
    if not os.path.exists(target_dir):
        os.makedirs(target_dir, exist_ok=True)

def is_valid_scene_name(scene_name):
    if scene_name.startswith('shovel_unloading_'):
        try:
            number = int(scene_name[len('shovel_unloading_'):])
            return 1 <= number <= 59
        except ValueError:
            return False
    return False


dir_current_file = os.path.dirname(__file__)  # 'Algorithm_1\planner'
dir_parent_1 = os.path.dirname(dir_current_file)  # 'Algorithm_1'


def plot_dubins_pose(observation, color='r', s=1, ax=None):
    if ax is None:
        ax = plt.gca()  # 如果未提供ax参数，则使用当前的Axes对象
    for i in range(len(observation["hdmaps_info"]["tgsc_map"].dubins_pose)):
        dubins_x = observation["hdmaps_info"]["tgsc_map"].dubins_pose[i]['x']
        dubins_y = observation["hdmaps_info"]["tgsc_map"].dubins_pose[i]['y']
        ax.scatter(dubins_x, dubins_y, color=color, s=s)


def plot_borderline(observation, linewidth=0.5, ax=None):
    if ax is None:
        ax = plt.gca()
    for i in range(len(observation["hdmaps_info"]["tgsc_map"].borderline)):
        borderpoints = observation["hdmaps_info"]["tgsc_map"].borderline[i]['borderpoints']
        x_values = [point[0] for point in borderpoints]
        y_values = [point[1] for point in borderpoints]
        ax.plot(x_values, y_values, linewidth=linewidth)


def plot_reference_path(observation, linewidth=0.5, ax=None):
    if ax is None:
        ax = plt.gca()
    for i in range(len(observation["hdmaps_info"]["tgsc_map"].reference_path)):
        waypoints = observation["hdmaps_info"]["tgsc_map"].reference_path[i]['waypoints']
        x_values = [point[0] for point in waypoints]
        y_values = [point[1] for point in waypoints]
        ax.plot(x_values, y_values, linestyle='--', linewidth=linewidth)


def plot_goal(goal_x, goal_y, ax=None):
    if ax is None:
        ax = plt.gca()
    polygon = [[goal_x[0], goal_y[0]], [goal_x[1], goal_y[1]], [goal_x[2], goal_y[2]], [goal_x[3], goal_y[3]]]
    ax.fill(*zip(*polygon), color='y', alpha=0.4, label='goal')


def get_utm_local_range(x_min, x_max, y_min, y_max, start_x, start_y, goal_x, goal_y):
    x_coords = ([start_x] + goal_x + [x_min, x_max])
    y_coords = ([start_y] + goal_y + [y_min, y_max])
    x_min, x_max = min(x_coords), max(x_coords)
    y_min, y_max = min(y_coords), max(y_coords)
    max_range = max(x_max - x_min, y_max - y_min)  # 用于确定正方形的边长
    x_center, y_center = (x_min + x_max) / 2, (y_min + y_max) / 2
    utm_local_range = (
        x_center - max_range / 2, y_center - max_range / 2,
        x_center + max_range / 2, y_center + max_range / 2
    )
    return utm_local_range


def expand_zeros(downsampled_image, start_ii, start_jj, a):
    rows = len(downsampled_image)
    cols = len(downsampled_image[0])
    expanded_image = [[1] * cols for _ in range(rows)]

    a1 = a*0.7    # 起点拓宽，但不要太宽
    # 遍历downsampled_image矩阵
    for i in range(rows):
        for j in range(cols):
            if downsampled_image[i][j] == 0:
                expanded_image[i][j] = 0
            else:
                # 检查当前位置的相邻元素是否为0
                if (i > 0 and downsampled_image[i - 1][j] == 0) or \
                        (i < rows - 1 and downsampled_image[i + 1][j] == 0) or \
                        (j > 0 and downsampled_image[i][j - 1] == 0) or \
                        (j < cols - 1 and downsampled_image[i][j + 1] == 0):
                    expanded_image[i][j] = 0
            if (start_jj-a1) <= i <= (start_jj + a1) and (start_ii-a1) <= j <= (start_ii + a1):
                expanded_image[i][j] = 1

    return np.array(expanded_image)


def downsample_image(image, target_shape):
    image_shape = image.shape
    downsampled_image_1 = np.zeros(target_shape, dtype=bool)

    row_ratio = image_shape[0] / target_shape[0]
    col_ratio = image_shape[1] / target_shape[1]

    for i in range(target_shape[0]):
        for j in range(target_shape[1]):
            original_i = int(i * row_ratio)
            original_j = int(j * col_ratio)
            downsampled_image_1[i, j] = image[original_i, original_j]

    return downsampled_image_1


def plot_bit_path(downsampled_image_2, bit_path):
    xx = np.arange(downsampled_image_2.shape[1])
    yy = np.arange(downsampled_image_2.shape[0])
    xx, yy = np.meshgrid(xx, yy)
    plt.pcolormesh(xx, yy, downsampled_image_2, cmap='gray')
    if isinstance(bit_path, list):
        x_coords = [point[0] for point in bit_path]
        y_coords = [point[1] for point in bit_path]
        plt.plot(x_coords, y_coords, linewidth=2, label='path', alpha=0.7)
    else:
        plt.scatter(*bit_path, s=6, label='path', alpha=0.8)
    plt.legend(loc='upper right')
    # plt.colorbar()  # 添加颜色条


def path_interpolation(point1, point2, max_distance):
    distance = math.sqrt((point2[0] - point1[0]) ** 2 + (point2[1] - point1[1]) ** 2)
    num_points = max(int(distance / max_distance), 1)
    x_step = (point2[0] - point1[0]) / num_points
    y_step = (point2[1] - point1[1]) / num_points

    interpolated_points = []
    for i in range(num_points + 1):
        x = point1[0] + i * x_step
        y = point1[1] + i * y_step
        interpolated_points.append((x, y))

    return interpolated_points


def interpolate_path(path, max_distance):
    interpolated_path = [path[0]]
    for i in range(len(path) - 1):
        point1 = path[i]
        point2 = path[i + 1]
        interpolated_points = path_interpolation(point1, point2, max_distance)
        interpolated_path.extend(interpolated_points[1:])

    return interpolated_path


class RRTNode:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.parent = None

def real2bit(real_x, real_y):
    closest_x = min(grid_map, key=lambda row: abs(row[0][0] - real_x))
    bit_x = grid_map.index(closest_x)
    closest_y = min(grid_map[bit_x], key=lambda val: abs(val[1] - real_y))
    bit_y = grid_map[bit_x].index(closest_y)
    return (bit_x, bit_y)

def bit2real(bit_x, bit_y, grid_map):
    real_x = x_min + (bit_x+1) * (x_max - x_min) / len(grid_map)
    real_y = y_min + (bit_y+1) * (y_max - y_min) / len(grid_map[0])

    return real_x, real_y

def calculate_path_length(bit_path):
    total_length = 0
    for i in range(len(bit_path) - 1):
        x1, y1 = bit_path[i]
        x2, y2 = bit_path[i + 1]
        distance = ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5
        total_length += distance
    return total_length

def back_cars2_obstacles_1(x_1, y_1, local_x_range, local_y_range, image_ndarray, value=False, distance=10):
    x_bit = (x_1 - local_x_range[0]) * 10
    y_bit = (y_1 - local_y_range[0]) * 10
    # 创建原始数组的副本
    image_ndarray_copy = np.copy(image_ndarray)
    for i in range((int(x_bit-distance)), int(x_bit+distance)):
        for j in range((int(y_bit-distance)), int(y_bit+distance)):
            indices = (int(j),int(i))
            image_ndarray_copy[indices] = value
    # 将对应索引位置的像素设置为False
    return image_ndarray_copy

def start2_obstacles_1(x_1, y_1, local_x_range, local_y_range, image_ndarray, value=True):
    x_bit = (x_1 - local_x_range[0]) * 10
    y_bit = (y_1 - local_y_range[0]) * 10
    # 创建原始数组的副本
    image_ndarray_copy = np.copy(image_ndarray)
    for i in range((int(x_bit)-30), int(x_bit+30)):
        for j in range((int(y_bit)-30), int(y_bit+30)):
            indices = (int(j),int(i))
            image_ndarray_copy[indices] = value
    # 将对应索引位置的像素设置为False
    return image_ndarray_copy

def add_trajectory_to_image_ndarray(observation, traj_true_dict,traj,image_ndarray_new):
    # image_ndarray_new = observation['hdmaps_info']['image_mask'].image_ndarray
    start_x = observation["vehicle_info"]["ego"]["x"]
    start_y = observation["vehicle_info"]["ego"]["y"]
    local_x_range = observation['hdmaps_info']['image_mask'].bitmap_info['bitmap_mask_PNG']['UTM_info']['local_x_range']
    local_y_range = observation['hdmaps_info']['image_mask'].bitmap_info['bitmap_mask_PNG']['UTM_info']['local_y_range']
    for i in range(len(traj_true_dict)):
        key = list(traj_true_dict.keys())[i]
        backvehicle_traj = traj_true_dict[key]
        if backvehicle_traj == []:
            continue
        if observation['test_setting']['scenario_type'] == "unloading":
            backvehicle_traj = backvehicle_traj[:50]
        elif observation['test_setting']['scenario_type'] == 'intersection':    # 交叉口动态规划
            # location = observation['hdmaps_info']['image_mask'].location
            # if location == 'guangdong_dapai':
            #     backvehicle_traj = backvehicle_traj[:50]
            # elif location == 'jiangxi_jiangtong':
            if scene_name in ['jiangtong_intersection_6_3_1', 'jiangtong_intersection_6_3_6',
                              'jiangtong_intersection_1_2_6', 'jiangtong_intersection_1_3_10', 'jiangtong_intersection_5_3_2'] :
                backvehicle_traj = backvehicle_traj[:50]
            if scene_name in ['jiangtong_intersection_1_3_2',
                              'jiangtong_intersection_1_2_5']:
                backvehicle_traj = backvehicle_traj[:50]

        vehicle_width = traj[key]['shape']['width']
        dist_1_start = math.sqrt((start_x - backvehicle_traj[0][0]) ** 2 + (start_y - backvehicle_traj[0][1]) ** 2) * 10

        xback = backvehicle_traj[-1][0]
        yback = backvehicle_traj[-1][1]
        image_ndarray_new = back_cars2_obstacles_1(xback, yback, local_x_range, local_y_range,
                                                   image_ndarray_new, value=False, distance=vehicle_width * 10)
        i = 0
        while True:
            xback = backvehicle_traj[i][0]
            yback = backvehicle_traj[i][1]
            image_ndarray_new = back_cars2_obstacles_1(xback, yback, local_x_range, local_y_range,
                                                       image_ndarray_new, value=False, distance=vehicle_width*10)
            i = i+16     # 隔几个点取一个
            if (i > (len(backvehicle_traj) - 1)):
                image_ndarray_new = start2_obstacles_1(start_x, start_y, local_x_range, local_y_range,
                                                       image_ndarray_new, value=True)
                break
    pixel_local_range = observation['hdmaps_info']['image_mask'].bitmap_local_info['pixel_local_range']
    return image_ndarray_new[pixel_local_range[1]:pixel_local_range[3], pixel_local_range[0]:pixel_local_range[2]]


class Planner_1:
    def __init__(self, grid_map, scenario_data, observation):
        self.grid_map = grid_map
        self.scenario_data = scenario_data
        self.goal_x = self.scenario_data["goal"]["x"]
        self.goal_y = self.scenario_data["goal"]["y"]
        self.observation = observation
        self.goal_x = self.scenario_data["goal"]["x"]
        self.goal_y = self.scenario_data["goal"]["y"]
        self.start_x = observation["vehicle_info"]["ego"]["x"]
        self.start_y = observation["vehicle_info"]["ego"]["y"]

    def plan(self, max_iterations, step_size, observation, traj_true_dict,traj,back_move=None):

        start_x = observation['vehicle_info']["ego"]["x"]
        start_y = observation['vehicle_info']["ego"]["y"]
        goal_x = observation['test_setting']["goal"]["x"]
        goal_y = observation['test_setting']["goal"]["y"]

        from a2_planner import back_cars2_obstacles,is_point_in_polygon
        image_ndarray_new = observation['hdmaps_info']['image_mask'].image_ndarray
        if 'vehicle_info' in observation and len(observation['vehicle_info'])>1:
            local_x_range = observation['hdmaps_info']['image_mask'].bitmap_info['bitmap_mask_PNG']['UTM_info'][
                'local_x_range']
            local_y_range = observation['hdmaps_info']['image_mask'].bitmap_info['bitmap_mask_PNG']['UTM_info'][
                'local_y_range']

            image_ndarray = observation['hdmaps_info']['image_mask'].image_ndarray

            image_ndarray_new = image_ndarray
            if observation["test_setting"]["scenario_type"] != "intersection":
                for i in range(len(observation['vehicle_info'])):
                    key = list(observation['vehicle_info'])[i]
                    if key == 'ego':
                        continue
                    x_1 = observation['vehicle_info'][key]['x']
                    y_1 = observation['vehicle_info'][key]['y']
                    dist_1_start = math.sqrt((start_x - x_1) ** 2 + (start_y - y_1) ** 2) * 10

                    if is_point_in_polygon(x_1, y_1, goal_x, goal_y):
                        image_ndarray_new = image_ndarray
                    else:
                        image_ndarray_new = back_cars2_obstacles(x_1, y_1, local_x_range, local_y_range, image_ndarray_new,
                                                                 dist_1_start=dist_1_start, value=False)

        (start_ii,start_jj) = real2bit(self.start_x, self.start_y)
        tree = [RRTNode(start_ii,start_jj)]
        bit_goal_x, bit_goal_y = [], []
        for idx in range(len(goal_x)):
            (idx_ii, idx_jj) = real2bit(goal_x[idx], goal_y[idx])
            bit_goal_x.append(idx_ii)
            bit_goal_y.append(idx_jj)
        image_ndarray_local = add_trajectory_to_image_ndarray(observation, traj_true_dict,traj,image_ndarray_new)
        downsampled_image_1 = downsample_image(image_ndarray_local, (grid_size, grid_size))
        # if not traj or not back_move:    # 没有背景车或背景车不动
        #     avoid_distance = 13  # 避障距离
        if observation["test_setting"]["scenario_type"] == "unloading":
            avoid_distance = 13  # 避障距离
        elif observation["test_setting"]["scenario_type"] == 'intersection':
            avoid_distance = 7  # 避障距离
            # if scene_name == 'jiangtong_intersection_1_3_2':
            #     avoid_distance = 7  # 避障距离
        num_expand = int(round(avoid_distance // math.sqrt(x_step ** 2 + y_step ** 2)))
        # ego_width = observation['vehicle_info']['ego']['shape']['width']
        # lim_num_expand = int(round(ego_width // math.sqrt(x_step ** 2 + y_step ** 2)))
        downsampled_image_2 = downsampled_image_1
        for _ in range(num_expand):
            downsampled_image_2 = expand_zeros(downsampled_image_2, start_ii,start_jj, num_expand)

        i_plan = 0
        while i_plan < max_iterations:
            i_plan += 1
            i = random.randint(0, len(self.grid_map) - 1)
            j = random.randint(0, len(self.grid_map[0]) - 1)
            rand_node = RRTNode(i, j)

            nearest_node = self.get_nearest_node(tree, rand_node)  # 在树中找到与采样点最近的节点
            new_node = self.generate_new_node(nearest_node, rand_node, step_size)  # 根据步长生成新节点
            i_new_node, j_new_node = int(round(new_node.x)), int(round(new_node.y))

            # if len(tree) == 1:
            #     angle_start = math.atan2((i_new_node - start_ii), (j_new_node - start_jj))
            #     steer_ini = calculate_angle_difference(angle_start, pi_2_pi(observation["vehicle_info"]["ego"]['yaw_rad']))
            #     if steer_ini > math.radians(lim_angle):
            #         print(f"Initial steer angle: {math.degrees(steer_ini):.1f}")
            #         continue
            look_up = downsampled_image_2[j_new_node, i_new_node]
            if look_up:
                tree.append(new_node)  # 将新节点添加到树中
                new_node.parent = nearest_node

            arrive = self.is_in_goal_region(new_node, bit_goal_x, bit_goal_y) and look_up
            if arrive is False and i_plan == max_iterations-1:
                if dynamic_plot == True:
                    plt.figure(figsize=(5, 5))
                    plot_bit_path(downsampled_image_2, [])
                    plt.show()
                # print(f"###log### No feasible path at {observation['test_setting']['t']:.1f}s, <num_expand-{num_expand}>")
                return None     # 不重新寻找
                # i_plan = 0  # 缩小边界扩宽的距离，重新寻找
                # num_expand -= 1
                # if num_expand <= lim_num_expand:
                #     plot_bit_path(downsampled_image_2, [])
                #     plt.show()
                #     return None
                # downsampled_image_2 = downsampled_image_1
                # for _ in range(num_expand):
                #     downsampled_image_2 = expand_zeros(downsampled_image_2, start_ii,start_jj, num_expand)
                # continue

            if arrive:  # 判断是否达到目标区域
                print(f"###log### Find way at {observation['test_setting']['t']:.1f}s, <num_expand-{num_expand}>")
                bit_path = self.extract_path(tree)
                for _ in range(20):
                    i_pop = 1
                    while i_pop < len(bit_path) - 1:

                        point1 = bit_path[i_pop - 1]
                        point2 = bit_path[i_pop + 1]

                        bit_path_pop = bit_path[i_pop]
                        bit_path.pop(i_pop) # 跳点

                        num_points = 200
                        x_interp = np.linspace(point1[0], point2[0], num_points)
                        y_interp = np.linspace(point1[1], point2[1], num_points)

                        collision_mat = np.zeros(x_interp.shape, dtype=bool)
                        for i_check in range(len(x_interp)):
                            x_check = int(round(x_interp[i_check]))
                            y_check = int(round(y_interp[i_check]))
                            collision_mat[i_check] = downsampled_image_2[y_check, x_check]

                        collision = np.any(collision_mat == False)
                        if collision:
                            bit_path.insert(i_pop, bit_path_pop)    # 撞了，不跳
                            i_pop += 1
                            continue

                        # if i_pop < len(bit_path) - 2:
                        #     point2 = bit_path[i_pop]   # 检查偏航角，大于阈值则不剔除i_pop
                        #     point3 = bit_path[i_pop + 1]    # 此处本来是加2，但是前面pop掉了一个
                        #     if i_pop == 1:
                        #         angle0 = pi_2_pi(observation["vehicle_info"]["ego"]['yaw_rad'])
                        #     else:
                        #         point0 = bit_path[i_pop - 2]
                        #         angle0 = math.atan2(point1[1] - point0[1], point1[0] - point0[0])
                        #     angle1 = math.atan2(point2[1] - point1[1], point2[0] - point1[0])
                        #     angle2 = math.atan2(point3[1] - point2[1], point3[0] - point2[0])
                        #     steer_angle0 = calculate_angle_difference(angle1, angle0)
                        #     steer_angle1 = calculate_angle_difference(angle2, angle1)
                        #     if steer_angle1 > math.radians(lim_angle) or steer_angle0 > math.radians(lim_angle):
                        #         bit_path.insert(i_pop, bit_path_pop)
                        #         i_pop += 1

                max_distance = calculate_path_length(bit_path)/100
                interp_bit_path = interpolate_path(bit_path, max_distance)

                for _ in range(40):
                    interp_bit_path = self.cute_path(interp_bit_path, downsampled_image_2, observation)

                real_path = [bit2real(*interp_bit_path[i], self.grid_map) for i in range(len(interp_bit_path))]

                file_path = os.path.join(dir_parent_1, "outputs", "plot", scene_name + "_bit" + "_time_" + str(round(observation['test_setting']['t'], 1)) + "s" +".png")
                os.makedirs(os.path.dirname(file_path), exist_ok=True)

                if do_plot == True:
                    plt.figure(figsize=(5, 5))
                    plot_bit_path(downsampled_image_2, interp_bit_path)
                    plt.scatter(start_ii,start_jj, marker='o', color='b', s=30, label='start', facecolors='white')
                    plot_goal(bit_goal_x, bit_goal_y)
                    plt.legend(loc='upper right')
                    plt.savefig(file_path)

                return real_path

    def cute_path(self, bit_path, downsampled_image_2, observation):
        i_smo = 1
        while i_smo < len(bit_path) - 1:

            point1 = bit_path[i_smo - 1]
            point2 = bit_path[i_smo + 1]

            num_points = 200
            x_interp = np.linspace(point1[0], point2[0], num_points)
            y_interp = np.linspace(point1[1], point2[1], num_points)

            collision_mat = np.zeros(x_interp.shape, dtype=bool)
            for i_check in range(len(x_interp)):
                x_check = int(round(x_interp[i_check]))
                y_check = int(round(y_interp[i_check]))
                collision_mat[i_check] = downsampled_image_2[y_check, x_check]

            collision = np.any(collision_mat == False)

            if not collision:
                bit_path_smo = bit_path[i_smo]
                bit_path[i_smo] = (x_interp[100], y_interp[100])

                # if i_smo < len(bit_path) - 2:
                #     if i_smo == 0:
                #         angle0 = pi_2_pi(observation["vehicle_info"]["ego"]['yaw_rad'])
                #     else:
                #         point0 = bit_path[i_smo - 2]
                #         angle0 = math.atan2(point1[1] - point0[1], point1[0] - point0[0])
                #     point3 = bit_path[i_smo + 2]
                #     angle1 = math.atan2(point2[1] - point1[1], point2[0] - point1[0])
                #     angle2 = math.atan2(point3[1] - point2[1], point3[0] - point2[0])
                #     steer_angle0 = calculate_angle_difference(angle1, angle0)
                #     steer_angle1 = calculate_angle_difference(angle2, angle1)  # 检查偏航角
                #     if steer_angle1 > math.radians(lim_angle) or steer_angle0 > math.radians(lim_angle):
                #         bit_path[i_smo] = bit_path_smo

            i_smo += 1

        return bit_path


    def get_nearest_node(self, tree, node):
        nearest_node = None
        min_dist = float('inf')

        for n in tree:
            dist = math.sqrt((n.x - node.x) ** 2 + (n.y - node.y) ** 2)
            if dist < min_dist:
                nearest_node = n
                min_dist = dist

        return nearest_node

    def generate_new_node(self, nearest_node, rand_node, step_size):
        dist = math.sqrt((rand_node.x - nearest_node.x) ** 2 + (rand_node.y - nearest_node.y) ** 2)

        if dist <= step_size:
            return rand_node

        theta = math.atan2(rand_node.y - nearest_node.y, rand_node.x - nearest_node.x)
        x = nearest_node.x + step_size * math.cos(theta)
        y = nearest_node.y + step_size * math.sin(theta)
        return RRTNode(x, y)

    def is_in_goal_region(self, node, goal_x, goal_y) -> bool:
        if node is None:
            return False

        goal_area_coords = list(zip(goal_x, goal_y))  # 在目标区域创建多边形
        goal_area_polygon = Polygon(goal_area_coords)

        if goal_area_polygon.contains(Point(node.x, node.y)):
            return True
        return False

    def extract_path(self, tree):
        path = []
        node = tree[-1]

        while node is not None:
            path.append((node.x, node.y))
            node = node.parent

        return path[::-1]


def calculate_yaw_angles(path):
    yaw_angles = []

    for i in range(len(path) - 1):
        dx = path[i + 1][0] - path[i][0]
        dy = path[i + 1][1] - path[i][1]
        yaw = math.atan2(dy, dx)
        yaw_angles.append(yaw)

    return yaw_angles

def calculate_yaw_differences(yaw_angles):
    yaw_differences = []
    for i in range(1, len(yaw_angles)):
        yaw_diff = calculate_angle_difference(yaw_angles[i], yaw_angles[i - 1])
        yaw_differences.append(yaw_diff)
    return yaw_differences

def calculate_angle_difference(angle1, angle2):

    difference = angle2 - angle1

    # 归一化到 [-π, π] 区间
    difference = (difference + math.pi) % (2 * math.pi) - math.pi

    return abs(difference)

def process(scenario_to_test, observation, traj_true_dict,traj,back_move=None):

    global scene_name
    scene_name = scenario_to_test['data']['scene_name']  # 获取场景名
    if scene_name in ['jiangtong_intersection_6_3_1'] and observation['test_setting']['t']>4:
        return None, []
    if scene_name in ['jiangtong_intersection_6_3_6'] and observation['test_setting']['t']>5:
        return None, []
    if scene_name in ['jiangtong_intersection_1_2_6'] and observation['test_setting']['t']>8:
        return None, []
    if scene_name in ['jiangtong_intersection_1_3_10'] and observation['test_setting']['t']>43:
        return None, []
    if scene_name in ['jiangtong_intersection_5_3_2'] and observation['test_setting']['t']>17:
        return None, []
    if scene_name == 'shovel_unloading_56':
        return None, []
    json_file = os.path.join(dir_parent_1, "inputs", "Scenarios", f"Scenario-{scene_name}.json")  # 构造JSON文件路径

    with open(json_file, "r") as f:  # 加载JSON文件
        scenario_data = json.load(f)

    global grid_size
    if observation["test_setting"]["scenario_type"] == "unloading":
        grid_size = 600  # 网格大小
    elif observation["test_setting"]["scenario_type"] == 'intersection':
        grid_size = 400  # 网格大小

    global grid_map, x_min, x_max, y_min, y_max, goal_x, goal_y, start_x, start_y, x_step, y_step
    x_min = scenario_data["x_min"]  # 提取需要的参数
    x_max = scenario_data["x_max"]
    y_min = scenario_data["y_min"]
    y_max = scenario_data["y_max"]
    max_t = scenario_data["max_t"]
    goal_x = scenario_data["goal"]["x"]
    goal_y = scenario_data["goal"]["y"]
    start_x = scenario_data["ego_info"]["start_states"]["x"]
    start_y = scenario_data["ego_info"]["start_states"]["y"]
    # start_x = observation["vehicle_info"]["ego"]["x"]
    # start_y = observation["vehicle_info"]["ego"]["x"]

    utm_local_range = observation['hdmaps_info']['image_mask'].bitmap_local_info['utm_local_range']
    x_min, y_min, x_max, y_max = utm_local_range[0], utm_local_range[1], utm_local_range[2], utm_local_range[3]
    x_margin, y_margin = 20, 20
    x_min, x_max, y_min, y_max = x_min - x_margin, x_max + x_margin, y_min - y_margin, y_max + y_margin

    grid_map = [[0 for _ in range(grid_size)] for _ in range(grid_size)]  # 初始化网格地图
    x_range = x_max - x_min  # X 范围
    y_range = y_max - y_min  # Y 范围
    x_step = x_range / grid_size  # X 步长
    y_step = y_range / grid_size  # Y 步长
    for i in range(grid_size):  # 根据全局坐标范围和步长填充网格地图
        for j in range(grid_size):
            x = x_min + i * x_step
            y = y_min + j * y_step
            grid_map[i][j] = (x, y)

    planner_1 = Planner_1(grid_map, scenario_data, observation)
    if observation['test_setting']['scenario_type'] == "unloading":
        real_step_size = 10  # 寻路步长
    else:
        real_step_size = 5  # 寻路步长
    bit_step_size = int(round(real_step_size // math.sqrt(x_step ** 2 + y_step ** 2)))
    if observation['test_setting']['scenario_type'] == "unloading":
        path = planner_1.plan(100000, bit_step_size, observation, traj_true_dict, traj)
    else:
        path = planner_1.plan(5000, bit_step_size, observation, traj_true_dict,traj)
    if path is not None:
        path = [list(coord) for coord in path]
        path = interpolate_path(path, 0.1)

        # 计算路径中每两个点之间的朝向角
        yaw_angles = calculate_yaw_angles(path)
        yaw_angles.insert(0, pi_2_pi(observation["vehicle_info"]["ego"]['yaw_rad']))

        # 计算朝向角列表中的元素两两相减，得到偏航角的列表
        yaw_differences = calculate_yaw_differences(yaw_angles)
        print(f"Max steer angle: {math.degrees(max(yaw_differences)):.1f}")
        location = observation['hdmaps_info']['image_mask'].location
        if location == 'guangdong_dapai':
            lim_angle = 30  # 最大偏航角
        elif location == 'jiangxi_jiangtong':
            lim_angle = 45  # 最大偏航角
            if scene_name == 'jiangtong_intersection_6_2_1':
                lim_angle = 100  # 最大偏航角
        if observation['test_setting']['scenario_type'] != "unloading" and max(yaw_differences) >= math.radians(lim_angle):
            return None, []
        # if (not traj or not back_move) and max(yaw_differences) >= math.radians(50):
        #     return None, []     # 如果没有背景车，最大转角50度

        path = [list(p) + [yaw] + [True] for p, yaw in zip(path, yaw_angles)]

        target_speed_f = 11
        spd_planned = (target_speed_f * np.ones(len(path), dtype=float)).tolist()
        # 计算前10米的点
        distance = 0
        points_5m = []
        for i in range(1, len(path)):
            dx = path[i][0] - path[i - 1][0]
            dy = path[i][1] - path[i - 1][1]
            distance += np.sqrt(dx ** 2 + dy ** 2)
            if distance <= 5:
                points_5m.append(i)
            else:
                break

        points_5m = [x-1 for x in points_5m]
        # 处理前10米的速度点
        if points_5m:
            start_speed = observation["vehicle_info"]["ego"]['v_mps']
            end_speed = target_speed_f

            # 线性插值
            for i, idx in enumerate(points_5m):
                spd_planned[idx] = start_speed + (end_speed - start_speed) * (i / (len(points_5m) - 1))

        # 更新spd_planned列表，确保第一个点和最后一个点的速度值
        if points_5m:
            spd_planned[points_5m[0]] = observation["vehicle_info"]["ego"]['v_mps']
            spd_planned[points_5m[-1]] = target_speed_f


        file_path = os.path.join(dir_parent_1, "outputs", "plot", scene_name +  "_real"  + "_time_" + str(round(observation['test_setting']['t'], 1)) + "s" + ".png")
        os.makedirs(os.path.dirname(file_path), exist_ok=True)

        if do_plot == True:
            plt.figure(figsize=(6.2, 5), dpi=400)

            if observation['test_setting']['scenario_type'] == "unloading":
                from a2_planner import plot_path_planned
                plot_path_planned(observation, path, start_x, start_y, goal_x, goal_y,
                                  [], x_min, x_max, y_min, y_max, [], [], spd_planned, scene_name)
            elif observation['test_setting']['scenario_type'] == "intersection":
                plot_dubins_pose(observation, color='r', s=1)
                plot_borderline(observation, linewidth=0.5)
                plot_reference_path(observation, linewidth=0.5)
                plt.scatter(start_x, start_y, marker='o', color='b', s=30, label='start', facecolors='white')
                plot_goal(goal_x, goal_y)
                plt.xlim(x_min, x_max)
                plt.ylim(y_min, y_max)
                x_coords = [point[0] for point in path]
                y_coords = [point[1] for point in path]
                plt.scatter(x_coords, y_coords, c=spd_planned, s=1, linewidth=2, label='path', alpha=0.6, cmap='rainbow')
                plt.legend(loc='upper right')
                plt.title(scene_name)
                plt.colorbar(label='Speed (m/s)')
            plt.savefig(file_path)

        from a2_planner import calculate_curvature
        for i in range(1, len(path)):
            path[i].append(
                calculate_curvature(path[i - 1][0], path[i - 1][1], path[i - 1][2],
                                    path[i][0], path[i][1], path[i][2]))
        return path[1:], spd_planned[1:]
    return None, []

